"""
Author: Morphlng
Date: 2024-04-06 17:01:45
LastEditTime: 2025-01-21 20:03:07
LastEditors: Morphlng
Description: Semantic observation of the kinematic features, mimicking highway-env.
FilePath: /DrivingGym/src/driving_gym/environment/agent/obs/kinematic_obs.py
"""

from __future__ import annotations

import gymnasium as gym
import numpy as np
import pandas as pd

from driving_gym.data.data_provider import Snapshot
from driving_gym.environment.agent.obs.base_obs import BaseObs
from driving_gym.environment.scenario.base_scenario import BaseScenario
from driving_gym.misc.util import norm_angle, override
from driving_gym.simulation.adapter_interface import AdapterInterface
from driving_gym.simulation.common.geom import Transform, Vector3D


class KinematicObs(BaseObs):

    AVAILABLE_FEATURES = [
        "presence",
        "x",
        "y",
        "vx",
        "vy",
        "heading",
        "cos_h",
        "sin_h",
        "lon_offset",
        "lat_offset",
        "angle_offset",
    ]

    def __init__(self, config: dict, adapter: AdapterInterface):
        super().__init__(config, adapter)

        self.actor_count: int = config.get("actor_count", 5)
        self.features: list[str] = config.get("features", KinematicObs.AVAILABLE_FEATURES)  # fmt: skip
        self.normalize: bool = config.get("normalize", True)
        self.absolute: bool = config.get("absolute", False)
        self.focus_ids: list[str] = config.get("focus_ids", None)
        self.ignore_ids: list[str] = config.get("ignore_ids", [])

    @override(BaseObs)
    def get_obs(self, snapshot: Snapshot, scenario: BaseScenario) -> np.ndarray:
        frame_info = snapshot.data.copy()

        # 1. Prepare valid informations
        ego_info = frame_info.pop(self.actor_id)
        focus_ids = self.focus_ids or frame_info.keys()
        valid_ids = set(focus_ids) - set(self.ignore_ids)
        frame_info = {id: frame_info[id] for id in frame_info if id in valid_ids}

        ego_transfrom = ego_info["transform"]
        ego_velocity = ego_info["velocity"]
        ego_location = ego_transfrom.location
        distances = {
            id: ego_location.distance(info["transform"].location)
            for id, info in frame_info.items()
        }
        nearest_actors = sorted(distances, key=distances.get)[: self.actor_count - 1]
        frame_info = {actor_id: frame_info[actor_id] for actor_id in nearest_actors}

        # 2. Get the observation
        semantic_infos = {
            "ego": self.get_actor_state(ego_transfrom, ego_velocity),
            "hero": None,
        }
        for actor_id in frame_info:
            actor_info = frame_info[actor_id]
            actor_transform = actor_info["transform"]
            actor_velocity = actor_info["velocity"]
            actor_state = self.get_actor_state(actor_transform, actor_velocity)
            if not self.absolute:
                actor_state.update(
                    self.relative_state(actor_transform, actor_velocity, ego_transfrom)
                )
            if actor_id.lower() in ["ego", "hero"]:
                semantic_infos["hero"] = actor_state
            else:
                semantic_infos[actor_id] = actor_state

        if semantic_infos["hero"] is None:
            semantic_infos.pop("hero")

        df = pd.DataFrame.from_dict(
            semantic_infos, orient="index", columns=self.features
        )
        if self.normalize:
            for feat in ["x", "y", "vx", "vy"]:
                feat_min = df[feat].min()
                feat_max = df[feat].max()
                denominator = feat_max - feat_min
                df[feat] = (df[feat] - feat_min) / (denominator + 1e-8)
        df = df[self.features]
        df.fillna(0, inplace=True)

        obs = df.to_numpy(dtype=np.float32)
        if obs.shape[0] < self.actor_count:
            pad_width = self.actor_count - obs.shape[0]
            pad_array = np.zeros((pad_width, len(self.features)), dtype=np.float32)
            obs = np.vstack((obs, pad_array))

        return obs

    @override(BaseObs)
    def get_observation_space(self) -> gym.Space:
        shape = (self.actor_count, len(self.features))
        return gym.spaces.Box(low=-np.inf, high=np.inf, shape=shape, dtype=np.float32)

    def get_actor_state(self, transform: Transform, velocity: Vector3D):
        """Get the state of an actor, including all of its kinematic features.

        Args:
            transform (Transform): The actor's transform.
            velocity (Vector3D): The actor's velocity.

        Returns:
            dict: A dictionary containing the actor's state information.
        """
        hd_map = self.adapter.get_map()
        wpt = hd_map.get_closest_lane_waypoint(transform.location)
        heading = np.deg2rad(norm_angle(transform.rotation.yaw))
        road_offset = wpt.transform.inverse_transform_location(transform.location)
        angle_offset = norm_angle(wpt.transform.rotation.yaw - transform.rotation.yaw)

        return {
            "presence": 1,
            "x": transform.location.x,
            "y": transform.location.y,
            "vx": velocity.x,
            "vy": velocity.y,
            "heading": heading,
            "cos_h": np.cos(heading),
            "sin_h": np.sin(heading),
            "lon_offset": road_offset.x,
            "lat_offset": road_offset.y,
            "angle_offset": np.deg2rad(angle_offset),
        }

    def relative_state(
        self,
        actor_transform: Transform,
        actor_velocity: Vector3D,
        ego_transform: Transform,
    ):
        """Get the relative state of an actor with respect to the ego vehicle.

        Args:
            actor_transform (Transform): The actor's transform.
            actor_velocity (Vector3D): The actor's velocity.
            ego_transform (Transform): The ego vehicle's transform.

        Returns:
            dict: A dictionary containing the actor's relative state information.
        """
        actor_location = actor_transform.location
        relative_location = ego_transform.inverse_transform_location(actor_location)
        relative_velocity = ego_transform.inverse_transform_location(actor_velocity + actor_location) - relative_location  # fmt: skip
        heading = norm_angle(actor_transform.rotation.yaw - ego_transform.rotation.yaw)
        heading = np.deg2rad(heading)
        return {
            "x": relative_location.x,
            "y": relative_location.y,
            "vx": relative_velocity.x,
            "vy": relative_velocity.y,
            "heading": heading,
            "cos_h": np.cos(heading),
            "sin_h": np.sin(heading),
        }
